#include "StdAfx.h"
#include ".\lsysga.h"

#include "MutationHelpers.h"
#include "CrossoverHelpers.h"

#include <map>

#define NUM_LEVELS 5

LsysGa::PopulationElement::PopulationElement()
{
	m_BranchRule.SetRandomly();
	int nNumArrEls = m_BranchRule.GetProbArrNumEnts(NUM_LEVELS);
	m_arrParamVectors = new ParamVector[nNumArrEls];

	m_BranchRule.RandomlySetParamArr(NUM_LEVELS, m_arrParamVectors);
}

LsysGa::PopulationElement::~PopulationElement()
{
	if(m_arrParamVectors)
	{
		delete [] m_arrParamVectors;
	}
	m_arrParamVectors = NULL;
}

const LsysGa::PopulationElement & LsysGa::PopulationElement::operator=(
	const LsysGa::PopulationElement &src)
{
	m_BranchRule = src.m_BranchRule;
	m_bFitnessDetermined = src.m_bFitnessDetermined;
	m_lfFitness = src.m_lfFitness;

	int nNumArrEls = m_BranchRule.GetProbArrNumEnts(NUM_LEVELS);
	for(int i = 0; i < nNumArrEls; ++i)
	{
		m_arrParamVectors[i][0] = src.m_arrParamVectors[i][0];
		m_arrParamVectors[i][1] = src.m_arrParamVectors[i][1];
		m_arrParamVectors[i][2] = src.m_arrParamVectors[i][2];
	}

	return *this;
}

void LsysGa::PopulationElement::Draw()
{
	m_BranchRule.Draw(NUM_LEVELS, m_arrParamVectors);
}

void LsysGa::PopulationElement::Creep(PopulationElement & Src, double lfCreepWidth)
{
	m_BranchRule = Src.m_BranchRule;
	m_BranchRule.Creep(lfCreepWidth);
	
	int nPopEls = m_BranchRule.GetProbArrNumEnts(NUM_LEVELS);
	for(int i = 0; i < nPopEls; ++i)
	{
		m_arrParamVectors[i][0] = 
			CreepVal(Src.m_arrParamVectors[i][0], lfCreepWidth, 1.0);
		m_arrParamVectors[i][1] = 
			CreepVal(Src.m_arrParamVectors[i][1], lfCreepWidth, 1.0);
		m_arrParamVectors[i][2] = 
			CreepVal(Src.m_arrParamVectors[i][2], lfCreepWidth, 1.0);
	}
}

void LsysGa::PopulationElement::CrossOver(PopulationElement & Src1, 
										  double lfWeight1, 
										  PopulationElement & Src2, 
										  double lfWeight2)
{
	double lfProb = lfWeight1/(lfWeight1+lfWeight2);
	m_BranchRule.Crossover(Src1.m_BranchRule, Src2.m_BranchRule, lfProb);
	int nPopEls = m_BranchRule.GetProbArrNumEnts(NUM_LEVELS);
	for(int i = 0; i < nPopEls; ++i)
	{
		//int nWhichSelect = (GenericMath::RandRange(0.0, 1.0) > lfProb);
		//m_arrParamVectors[i][0] = Src1.m_arrParamVectors[i][0]+
		//	(nWhichSelect)*(Src2.m_arrParamVectors[i][0]-Src1.m_arrParamVectors[i][0]);
		//m_arrParamVectors[i][1] = Src1.m_arrParamVectors[i][1]+
		//	(nWhichSelect)*(Src2.m_arrParamVectors[i][1]-Src1.m_arrParamVectors[i][1]);
		//m_arrParamVectors[i][2] = Src1.m_arrParamVectors[i][2]+
		//	(nWhichSelect)*(Src2.m_arrParamVectors[i][2]-Src1.m_arrParamVectors[i][2]);
		m_arrParamVectors[i][0] = 
			ProbSelect(lfProb, Src1.m_arrParamVectors[i][0], Src2.m_arrParamVectors[i][0]);
		m_arrParamVectors[i][1] = 
			ProbSelect(lfProb, Src1.m_arrParamVectors[i][1], Src2.m_arrParamVectors[i][1]);
		m_arrParamVectors[i][2] = 
			ProbSelect(lfProb, Src1.m_arrParamVectors[i][2], Src2.m_arrParamVectors[i][2]);
	}
}

void LsysGa::InitArrays()
{
	m_arrCurrentPop = new PopulationElement[m_nPopSize];
	m_arrNextPop= new PopulationElement[m_nPopSize];
	m_arrCrossoverResults= new PopulationElement[m_nNumOffspring];
	m_arrMutationResults= new PopulationElement[m_nNumMutants];
}

LsysGa::LsysGa(void)
{
	m_pDrawWidget = NULL;
	SetPopParams(50, 25, 150);
	InitArrays();
	m_pImgCompare = NULL;
	m_pTmpImage = NULL;
}

LsysGa::LsysGa(int nPopSize, int nMutants, int nOffspring)
{
	m_pDrawWidget = NULL;
	SetPopParams(nPopSize, nMutants, nOffspring);
	InitArrays();
	m_pImgCompare = NULL;
	m_pTmpImage = NULL;
}

LsysGa::~LsysGa(void)
{
	if(m_arrCurrentPop)
	{
		delete [] m_arrCurrentPop;
	}
	m_arrCurrentPop = NULL;

	if(m_arrNextPop)
	{
		delete [] m_arrNextPop;
	}
	m_arrNextPop = NULL;

	if(m_arrCrossoverResults)
	{
		delete [] m_arrCrossoverResults;
	}
	m_arrCrossoverResults = NULL;

	if(m_arrMutationResults)
	{
		delete [] m_arrMutationResults;
	}
	m_arrMutationResults = NULL;
}

void LsysGa::SetDrawWidget(CDrawableWidget * pWidget) 
{ 
	m_pDrawWidget = pWidget; 
	m_pDrawWidget->AddCustomDraw(&m_CurrBest);
}

#define USE_L1_NORM 0

#define USE_EDGE_FOR_EVAL 0

double LsysGa::Evaluate(PopulationElement * pEvalMe)
{
	if(!pEvalMe)
	{
		return -1.0;
	}

	if(!m_pImgCompare || !m_pDrawWidget)
	{
		pEvalMe->m_lfFitness = 0.0;
		return 0.0;
	}


	

	int nSize = m_pDrawWidget->PopCustomDraw();
	ASSERT(nSize > 0);

	m_pDrawWidget->AddCustomDraw(pEvalMe);

	m_pDrawWidget->DrawScreen();

	m_pDrawWidget->GrabScreen(*m_pTmpImage);

	int nWidth = m_pImgCompare->GetWidth();
	int nHeight = m_pImgCompare->GetHeight();

	double lfResult = 0.0;

	// do stuff
	static const double lfErrorScale = 1.0/256.0;

	unsigned char * pBitsImageCompare = 
		static_cast<unsigned char *>(m_pImgCompare->GetBits());

	unsigned char * pBitsImageTmp = 
		static_cast<unsigned char *>(m_pTmpImage->GetBits());

	int nBytes = nHeight*nWidth*3;



	pBitsImageTmp -= (nBytes-nWidth*3);
	pBitsImageCompare -= (nBytes-nWidth*3);

	UINT64 nSum = 0;

	bool bIsTmpDibSection = m_pTmpImage->IsDIBSection();

#if USE_EDGE_FOR_EVAL
	nHeight--;
	nWidth--;
	for(int i = 0; i < nHeight; ++i)
	{
		for(int j = 0; j < nWidth; ++j)
		{
			int nEdge1 = (*pBitsImageCompare-pBitsImageCompare[3*nWidth]);
			nEdge1 *= nEdge1;
			int nTmp = (*pBitsImageCompare-pBitsImageCompare[3]);
			nEdge1 += nTmp*nTmp;

			int nEdge2 = (*pBitsImageTmp-pBitsImageTmp[3*nWidth]);
			nEdge2 *= nEdge2;
			nTmp = (*pBitsImageTmp-pBitsImageTmp[3]);
			nEdge2 += nTmp*nTmp;

			nSum += abs(nEdge2-nEdge1);

			++pBitsImageCompare;
			++pBitsImageTmp;
		}
		++pBitsImageCompare;
		++pBitsImageTmp;
	}
#else
	for(int i = 0; i < nBytes; ++i)
	{
		int nToAdd = *pBitsImageCompare;
		nToAdd -= *pBitsImageTmp;

#if USE_L1_NORM
		nToAdd = abs(nToAdd);
#else
		nToAdd *= nToAdd;
#endif
		
		nSum += nToAdd;
		++pBitsImageCompare;
		++pBitsImageTmp;
	}
#endif


#if USE_L1_NORM
	lfResult = static_cast<double>(nSum)/(256.0*nBytes*3);
#else
#if USE_EDGE_FOR_EVAL
	lfResult = static_cast<double>(nSum)/(2.0*256.0*256.0*nBytes*3);
#else
	lfResult = static_cast<double>(nSum)/(256.0*256.0*nBytes*3);
#endif
#endif

	

	pEvalMe->m_lfFitness = lfResult;

	return pEvalMe->m_lfFitness;
}

void LsysGa::EvaluatePopulation(PopulationElement * arrPopElems, int nSize)
{
	for(int i = 0; i < nSize; ++i)
	{
		Evaluate(&(arrPopElems[i]));
	}
}

struct gtDouble
{
  bool operator()(const double lf1, const double lf2) const
  {
    return (lf1-lf2 < 0);
  }
};

void LsysGa::SelectBest(PopulationElement * pArr1, int nSize1, 
						PopulationElement * pArr2, int nSize2)
{
	std::multimap<double, PopulationElement *, gtDouble> mapPopEls;

	
	for(int i = 0; i < nSize1; ++i)
	{
		mapPopEls.insert(
			std::pair<double, PopulationElement *>(pArr1[i].m_lfFitness, pArr1+i));
	}

	for(int i = 0; i < nSize2; ++i)
	{
		mapPopEls.insert(
			std::pair<double, PopulationElement *>(pArr2[i].m_lfFitness, pArr2+i));
	}

	
	std::multimap<double, PopulationElement *, gtDouble>::iterator Iter = 
		mapPopEls.begin();

	m_CurrBest = *((*Iter).second);
	for(int i = 0; i < m_nPopSize; ++i)
	{
		if(Iter == mapPopEls.end())
		{
			Iter = mapPopEls.begin();
		}
		m_arrNextPop[i] = *((*Iter).second);
		++Iter;
	}

	PopulationElement * pTmp = m_arrCurrentPop;
	m_arrCurrentPop = m_arrNextPop;
	m_arrNextPop = pTmp;

	static int nIter = 0;
	++nIter;
	TRACE("%d\t\t%lf\n", nIter, m_CurrBest.m_lfFitness);
}

void LsysGa::DoCrossOver()
{
	// do stuff
	for(int i = 0; i < m_nNumOffspring; ++i)
	{
		int nParent1 = rand()%m_nPopSize;
		int nParent2 = (rand()%(m_nPopSize-1));
		nParent2= (nParent1+nParent2+1)%m_nPopSize;

		m_arrCrossoverResults[i].CrossOver(
			m_arrCurrentPop[nParent1], 1.0-m_arrCurrentPop[nParent1].m_lfFitness,
			m_arrCurrentPop[nParent2], 1.0-m_arrCurrentPop[nParent2].m_lfFitness);
	}
	EvaluatePopulation(m_arrCrossoverResults, m_nNumOffspring);
	SelectBest(m_arrCurrentPop, m_nPopSize/5, m_arrCrossoverResults, m_nNumOffspring);
}

void LsysGa::DoCreep()
{
	// do stuff
	for(int i = 0; i < m_nNumMutants; ++i)
	{
		int nToMutate = rand()%m_nPopSize;

		m_arrMutationResults[i].Creep(m_arrCurrentPop[nToMutate], 0.1); // make it decay later
		
	}
	EvaluatePopulation(m_arrMutationResults, m_nNumMutants);
	SelectBest(m_arrCurrentPop, (m_nPopSize)/2, m_arrMutationResults, m_nNumMutants);
}

void LsysGa::IterateOneGeneration()
{
	if(!m_pDrawWidget)
	{
		return;
	}

	

	
	// do stuff
	DoCrossOver();

	
	DoCreep();
	

	int nSize = m_pDrawWidget->PopCustomDraw();
	ASSERT(nSize > 0);
	m_pDrawWidget->AddCustomDraw(&m_CurrBest);

}

void LsysGa::Train(int nGenerations)
{
	if(!m_pDrawWidget)
	{
		return;
	}
	CDrawableWidget * pPopped = m_pDrawWidget->PopLink();

	if(!m_pTmpImage)
	{
		m_pTmpImage = m_pDrawWidget->GetBlankImage(1.0, false);
	}


	EvaluatePopulation(m_arrCurrentPop, m_nPopSize);
	// do stuff
	for(int nIter = 0; nIter < nGenerations; ++nIter)
	{
		IterateOneGeneration();
	}

	if(pPopped)
	{
		m_pDrawWidget->LinkTo(*pPopped);
	}

	m_pDrawWidget->DrawScreen();

	double lfError = m_CurrBest.m_lfFitness;
	CString strErrorMetrics;
	strErrorMetrics.Format("Current best fitness is %lf", lfError);
	AfxMessageBox(strErrorMetrics);

	//if(m_pTmpImage)
	//{
	//	
	//}
	//m_pTmpImage = NULL;
}
